import os
import json

from langchain_community.embeddings.dashscope import DashScopeEmbeddings
from langchain_community.llms.tongyi import Tongyi
# 新的导入方式
from langchain_chroma import Chroma

from langchain_community.document_loaders import SQLDatabaseLoader
from langchain_community.utilities.sql_database import SQLDatabase

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.prompts import PromptTemplate
from dotenv import load_dotenv

load_dotenv()

# 连接数据库
db = SQLDatabase.from_uri('mysql+pymysql://root:123456@127.0.0.1:3306/p5')


# 封装基本功能类
class Exercise:

    # 初始化一些调用对象
    def __init__(self):
        # 1.构建向量数据库
        self.embeddings = DashScopeEmbeddings(
            model='text-embedding-v2',
            dashscope_api_key=os.getenv("DASHSCOPE_API_KEY")
        )
        # 2.构建切片器
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=20,
            chunk_overlap=10
        )

        # init初始化方法--->为了一运行就走这个函数中的东西
        # 将数据库提前创建
        self.get_retriever()

    # 处理文件小工具
    def get_documents(self):
        # 加载数据库数据
        loader = SQLDatabaseLoader(
            db=db,
            query='SELECT * FROM tb_exercise'
        )
        # 加载数据
        articles = loader.load()

        # 验证数据不为空
        if not articles:
            print("警告: 数据库中没有找到文章数据")
            articles = []

        # 过滤掉空内容或None的文档
        valid_articles = []
        for article in articles:
            if article.page_content and article.page_content.strip():
                valid_articles.append(article)

        if not valid_articles:
            print("警告: 没有有效的内容")
            # 创建一个默认文档避免错误
            from langchain.schema import Document
            valid_articles = [Document(page_content="默认内容", metadata={})]

        # 切片文件
        articles_split = self.text_splitter.split_documents(valid_articles)
        return articles_split

    # 向量化存储
    def set_vector(self):
        documents = self.get_documents()
        print(f"健身类 - 文档数量: {len(documents)}")

        try:
            vector = Chroma.from_documents(
                documents=documents,
                embedding=self.embeddings,
                persist_directory='chroma/report_plan',
                collection_name='stories',
            )
            print("Story类 - 向量数据库创建成功")
        except Exception as e:
            print(f"Story类 - 创建向量数据库时出错: {e}")
            # 尝试使用更简单的模型
            embeddings_v2 = DashScopeEmbeddings(
                model='text-embedding-v2',
                dashscope_api_key=os.getenv("DASHSCOPE_API_KEY")
            )
            vector = Chroma.from_documents(
                documents=documents,
                embedding=embeddings_v2,
                persist_directory='chroma/report_plan',
                collection_name='stories',
            )

        # 还要创建检索器
        self.retriever = vector.as_retriever(
            search_kwargs={'k': 3}
        )

    # 检索器-->进行检索的过程，顺便判断有没有向量库
    def get_retriever(self):
        # chroma是创建库+存储
        vector = Chroma(
            persist_directory='chroma/report_plan',
            embedding_function=self.embeddings,
            collection_name='stories',
        )
        # 获取一条数据，看看有没有数据(有没有库)，没有就创建，有就直接返回检索链
        try:
            # 检查向量数据库是否为空
            count = vector._collection.count()
            if count == 0:
                print("向量数据库为空，开始创建...")
                self.set_vector()  # 调用上面创建库的方法
            else:
                print(f"向量数据库已有 {count} 个文档，直接创建检索器")
                # 有的话直接创建检索链
                self.retriever = vector.as_retriever(
                    search_kwargs={'k': 3}
                )
        except Exception as e:
            print(f"检查向量数据库时出错: {e}，重新创建...")
            self.set_vector()


exercise = Exercise()